# coding=utf-8
import random
import cv2
import numpy as np

def add_alpha_channel(img):
    """ 为jpg图像添加alpha通道 """
    r_channel, g_channel, b_channel = cv2.split(img) # 剥离jpg图像通道
    alpha_channel = np.ones(b_channel.shape, dtype=b_channel.dtype) * 255 # 创建Alpha通道
    img_new = cv2.merge((r_channel, g_channel, b_channel, alpha_channel)) # 融合通道
    return img_new

def handel_img(img):
    imgGray = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY) # 转灰度图
    imgBlur = cv2.GaussianBlur(imgGray, (5, 5), 1) # 高斯模糊
    imgCanny = cv2.Canny(imgBlur, 60, 60) # Canny算子边缘检测
    return imgCanny

def conmatch(target, template):# 目标，模板
    # 判断jpg图像是否已经为4通道
    if target.shape[2] == 3:
        target = add_alpha_channel(target)
    # 判断jpg图像是否已经为4通道
    if template.shape[2] == 3:
        template = add_alpha_channel(template)
    img = handel_img(target)
    small_img = handel_img(template)
    res_TM_CCOEFF_NORMED = cv2.matchTemplate(img, small_img, 3)  
    (minVal, maxVal, minLoc, maxLoc) = cv2.minMaxLoc(res_TM_CCOEFF_NORMED)
    # 确定起点和终点的（x，y）坐标边界框
    (startX, startY) = maxLoc
    endX = startX + template.shape[1]
    endY = startY + template.shape[0]
    return startX, startY, endX, endY

def get_tracks(dis):
    v = 0
    t = random.uniform(0.3,0.9)
    tracks = []
    current = 0
    while current <= dis:
        mid = current / dis
        if mid < random.uniform(0.095,0.1):
            a = random.uniform(2,2.5)
        elif random.uniform(0.1,0.105) <= mid < random.uniform(0.45,0.455):
            a = random.uniform(5,5.5)
        elif random.uniform(0.455,0.46) <= mid < random.uniform(0.9,0.905):
            a = random.uniform(7,7.5)
        else:
            a = random.uniform(-0.55,-0.5)
        v0 = v
        s = abs(v0*t+0.5*a*(t**2))
        current += s
        tracks.append(round(s))
        v = v0+a*t
    tracks.append(dis-current)
    return tracks

# 直接执行脚本，语句之前和之后的代码都被执行，引用脚本，之前的语句被执行，之后的没有被执行
if __name__ == "__main__":
    # 2. 对比图片，计算距离
    target_path = 'target.png' # 目标文件路径
    template_path = 'template.png' # 模板文件路径
    # 读取图像
    target = cv2.imread(target_path, cv2.IMREAD_UNCHANGED)
    template = cv2.imread(template_path, cv2.IMREAD_UNCHANGED)
    # 匹配图像
    startX1, startY1, endX1, endY1 = conmatch(target, template)
    # 切割出可能存在缺口的部分
    target_cut = target[startY1:endY1, endX1:target.shape[1]]
    # 匹配图像
    startX2, startY2, endX2, endY2 = conmatch(target_cut, template)
    startX2 = startX2 + endX1 + int(template.shape[1]/2)
    startY2 = startY2 + startY1
    endX2 = startX2 + int(template.shape[1])
    endY2 = startY2 + int(template.shape[0])
    # 在图像上绘制边框
    cv2.rectangle(target, (startX2, startY2), (endX2, endY2), (255, 0, 0), 3)
    # 显示输出图像
    cv2.imshow("Output", target)
    cv2.waitKey(0)